#pragma once

#include <cstdio>
#include <thread>
#include <atomic>
#include <iostream>

extern "C" {
#include <libavcodec/codec.h>
#include <libavformat/avformat.h>
}

#include "ThreadSafeQueue.h"


class Demuxer {
public:
    Demuxer(const std::string& filename, ThreadSafeQueue<AVPacket*>& packet_queue, std::atomic<bool>& demux_done)
        : filename(filename), packet_queue(packet_queue), quit(false), demux_done(demux_done) {}

    bool init() {
        avformat_open_input(&fmt_ctx, filename.c_str(), nullptr, nullptr);
        avformat_find_stream_info(fmt_ctx, nullptr);

        for (unsigned int i = 0; i < fmt_ctx->nb_streams; i++) {
            if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
                vs_idx = i;
                codec_paras = fmt_ctx->streams[i]->codecpar;
                break;
            }
        }

        return vs_idx >= 0;
    }

    AVCodecParameters* get_codec_paras() {
        return codec_paras;
    }

    int get_width() {
        return codec_paras->width;
    }

    int get_height() {
        return codec_paras->height;
    }

    void start() {
        demux_thread = std::thread(&Demuxer::demux, this);
    }

    void stop() {
        std::printf("Try to stop demux thread\n");
        std::fflush(stdout);
        quit = true;
        if (demux_thread.joinable()) {
            demux_thread.join();
        }
        std::printf("Demux thread stopped\n");
    }

    AVFormatContext* get_fmt_ctx() {
        return fmt_ctx;
    }

    int get_video_stream_index() {
        return vs_idx;
    }
    
    AVStream* get_video_stream() const {
        return fmt_ctx->streams[vs_idx];
    }

    AVStream* get_audio_stream() const {
        return fmt_ctx->streams[as_idx];
    }

    ~Demuxer() {
        if (fmt_ctx) {
            avformat_close_input(&fmt_ctx);
        }
        // if (codec_paras) {
        //     avcodec_parameters_free(&codec_paras);
        // }
        
        std::printf("Demuxer destroyed\n");
    }
private:
    void demux() {
        while (!quit) {
            AVPacket* packet = av_packet_alloc();
            if (av_read_frame(fmt_ctx, packet) >= 0) {
                if (packet->stream_index == vs_idx) {
                    packet_queue.push(packet);
                } else {
                    av_packet_free(&packet);
                }
            } else {
                av_packet_free(&packet);
                break;
            }
        }
        demux_done = true;
        std::printf("Demux done\n");
        std::fflush(stdout);
    }

    std::string filename;
    AVFormatContext* fmt_ctx = nullptr;
    AVCodecParameters* codec_paras = nullptr;
    int vs_idx = -1;
    int as_idx = -1;
    ThreadSafeQueue<AVPacket*>& packet_queue;
    std::atomic<bool> quit;
    std::thread demux_thread;
    std::atomic<bool>& demux_done;
};
